"""
边侧节点通信客户端
专门负责与边侧节点的HTTP通信和数据交换
"""

import asyncio
from typing import Dict, List, Optional, Tuple, Any
from datetime import datetime

from database import Node
from utils import LogManager, HTTPClient
from exceptions import ValidationException, NetworkException
from .edge_types import (
    EdgeHealthStatus, EdgeSampleInfo, EdgeAggregationRequest, EdgeAggregationResponse,
    EdgeUploadStatus, EdgeModelDeployment, EdgeModelInfo, EdgeInferenceResult,
    EdgeSystemStatus, EdgeOperationResult, EdgeNodeStatus
)


class EdgeClient:
    """边侧节点通信客户端"""
    
    def __init__(self, http_timeout: int = 10, max_retries: int = 1):  # 缩短超时时间，减少重试次数
        self.logger = LogManager.get_logger(self.__class__.__name__)
        self.http_timeout = http_timeout
        self.max_retries = max_retries
    
    def _get_http_client(self, node: Node) -> HTTPClient:
        """获取边侧节点HTTP客户端"""
        if node.node_type != 'edge':
            raise ValidationException(f"节点 {node.node_id} 不是边侧节点")
        
        # 构建边侧节点URL
        edge_url = f"http://{node.service_ip}"
        if ':' not in (node.service_ip or ""):
            edge_url = f"http://{node.service_ip}:8001"
        
        return HTTPClient(edge_url, self.http_timeout, self.max_retries)
    
    def _is_connection_error(self, error_str: str) -> bool:
        """判断是否为连接错误"""
        connection_errors = [
            "timed out",
            "Connection refused", 
            "Connection timeout",
            "ConnectTimeoutError",
            "NameResolutionError",
            "Network unreachable",
            "No route to host"
        ]
        return any(err in error_str for err in connection_errors)
    
    async def check_health(self, node: Node) -> EdgeHealthStatus:
        """检查边侧节点健康状态"""
        try:
            client = self._get_http_client(node)
            success, result = client.get("health")
            
            if success:
                self.logger.debug(f"边侧节点健康检查成功: {node.node_id}")
                return EdgeHealthStatus(
                    node_id=node.node_id,
                    status=EdgeNodeStatus.HEALTHY,
                    details=result,
                    timestamp=datetime.now().isoformat()
                )
            else:
                error_msg = str(result.get("error", "未知错误"))
                self.logger.warning(f"边侧节点健康检查失败: {node.node_id}, {error_msg}")
                
                # 根据错误类型设置状态
                status = EdgeNodeStatus.OFFLINE if self._is_connection_error(error_msg) else EdgeNodeStatus.UNHEALTHY
                
                return EdgeHealthStatus(
                    node_id=node.node_id,
                    status=status,
                    error=error_msg,
                    timestamp=datetime.now().isoformat()
                )
                
        except Exception as e:
            error_msg = str(e)
            self.logger.error(f"边侧节点健康检查异常: {node.node_id}, {error_msg}")
            
            # 连接错误时设置为离线状态
            status = EdgeNodeStatus.OFFLINE if self._is_connection_error(error_msg) else EdgeNodeStatus.ERROR
            
            return EdgeHealthStatus(
                node_id=node.node_id,
                status=status,
                error=error_msg,
                timestamp=datetime.now().isoformat()
            )
    
    async def get_samples_info(self, node: Node) -> EdgeSampleInfo:
        """获取边侧节点样本信息"""
        try:
            client = self._get_http_client(node)
            success, result = client.get("samples/info")
            
            if success:
                self.logger.debug(f"获取边侧节点样本信息成功: {node.node_id}")
                return EdgeSampleInfo(
                    node_id=result.get("node_id", node.node_id),
                    task_name=result.get("task_name", ""),
                    collected_sample_count=result.get("collected_sample_count", 0),
                    uncollected_sample_count=result.get("uncollected_sample_count", 0),
                    uninferred_sample_count=result.get("uninferred_sample_count", 0),
                    total_sample_count=result.get("total_sample_count", 0),
                    difficulty_threshold_stats=result.get("difficulty_threshold_stats", {}),
                    timestamp=result.get("timestamp", datetime.now().isoformat())
                )
            else:
                error_msg = str(result.get("error", "获取样本信息失败"))
                if self._is_connection_error(error_msg):
                    raise NetworkException(f"节点 {node.node_id} 连接失败，节点可能离线")
                else:
                    raise NetworkException(f"获取样本信息失败: {error_msg}")
                
        except Exception as e:
            error_msg = str(e)
            self.logger.error(f"获取边侧节点样本信息异常: {node.node_id}, {error_msg}")
            if self._is_connection_error(error_msg):
                raise NetworkException(f"节点 {node.node_id} 连接失败，节点可能离线")
            else:
                raise
    
    async def request_samples_aggregation(
        self, 
        node: Node, 
        request: EdgeAggregationRequest
    ) -> EdgeAggregationResponse:
        """请求边侧节点进行样本聚合"""
        try:
            client = self._get_http_client(node)
            
            # 构建请求数据
            request_data = request.model_dump(exclude_none=True)
            
            success, result = client.post("samples/aggregate", json_data=request_data)
            
            if success:
                self.logger.debug(f"请求边侧节点样本聚合成功: {node.node_id}")
                return EdgeAggregationResponse(
                    success=result.get("success", True),
                    node_id=node.node_id,
                    message=result.get("message", ""),
                    selected_samples=result.get("selected_samples", []),
                    sample_count=result.get("sample_count", 0),
                    upload_info=result.get("upload_info"),
                    request_params=request_data,
                    timestamp=datetime.now().isoformat()
                )
            else:
                error_msg = str(result.get("error", "样本聚合请求失败"))
                if self._is_connection_error(error_msg):
                    raise NetworkException(f"节点 {node.node_id} 连接失败，节点可能离线")
                else:
                    raise NetworkException(f"样本聚合请求失败: {error_msg}")
                
        except Exception as e:
            error_msg = str(e)
            self.logger.error(f"请求边侧节点样本聚合异常: {node.node_id}, {error_msg}")
            if self._is_connection_error(error_msg):
                raise NetworkException(f"节点 {node.node_id} 连接失败，节点可能离线")
            else:
                raise
    
    async def get_upload_status(self, node: Node, task_id: str) -> EdgeUploadStatus:
        """获取边侧节点上传任务状态"""
        try:
            client = self._get_http_client(node)
            success, result = client.get(f"samples/upload/{task_id}")
            
            if success:
                self.logger.debug(f"获取边侧节点上传状态成功: {node.node_id}, 任务: {task_id}")
                return EdgeUploadStatus(
                    node_id=node.node_id,
                    task_id=task_id,
                    status=result.get("status", "unknown"),
                    details=result,
                    timestamp=datetime.now().isoformat()
                )
            else:
                raise NetworkException(f"获取上传状态失败: {result}")
                
        except Exception as e:
            self.logger.error(f"获取边侧节点上传状态异常: {node.node_id}, {str(e)}")
            raise
    
    async def deploy_model(self, node: Node, deployment: EdgeModelDeployment) -> EdgeOperationResult:
        """部署模型到边侧节点"""
        try:
            client = self._get_http_client(node)
            
            # 构建模型更新请求
            update_request = {
                "model_id": deployment.model_id,
                "model_url": deployment.model_url,
                "model_path": deployment.model_path
            }
            
            success, result = client.post("model/update", json_data=update_request)
            
            if success:
                self.logger.info(f"模型部署成功: {deployment.model_id} -> {node.node_id}")
                return EdgeOperationResult(
                    success=True,
                    node_id=node.node_id,
                    operation="deploy_model",
                    data=result,
                    timestamp=datetime.now().isoformat()
                )
            else:
                error_msg = str(result.get("error", "模型部署失败"))
                if self._is_connection_error(error_msg):
                    raise NetworkException(f"节点 {node.node_id} 连接失败，节点可能离线")
                else:
                    raise NetworkException(f"模型部署失败: {error_msg}")
                
        except Exception as e:
            error_msg = str(e)
            self.logger.error(f"部署模型到边侧节点异常: {node.node_id}, {error_msg}")
            if self._is_connection_error(error_msg):
                raise NetworkException(f"节点 {node.node_id} 连接失败，节点可能离线")
            else:
                raise
    
    async def get_model_info(self, node: Node) -> EdgeModelInfo:
        """获取边侧节点模型信息"""
        try:
            client = self._get_http_client(node)
            success, result = client.get("model/info")
            
            if success:
                self.logger.debug(f"获取边侧节点模型信息成功: {node.node_id}")
                return EdgeModelInfo(
                    node_id=node.node_id,
                    model_loaded=result.get("model_loaded", False),
                    current_model_id=result.get("current_model_id"),
                    model_path=result.get("model_path"),
                    model_metadata=result.get("model_metadata"),
                    timestamp=datetime.now().isoformat()
                )
            else:
                raise NetworkException(f"获取模型信息失败: {result}")
                
        except Exception as e:
            self.logger.error(f"获取边侧节点模型信息异常: {node.node_id}, {str(e)}")
            raise
    
    async def trigger_inference(self, node: Node) -> EdgeInferenceResult:
        """触发边侧节点推理"""
        try:
            client = self._get_http_client(node)
            success, result = client.post("inference/trigger")
            
            if success:
                self.logger.debug(f"触发边侧节点推理成功: {node.node_id}")
                return EdgeInferenceResult(
                    node_id=node.node_id,
                    success=result.get("success", True),
                    processed_samples=result.get("processed_samples", 0),
                    inference_time=result.get("inference_time"),
                    timestamp=datetime.now().isoformat()
                )
            else:
                raise NetworkException(f"触发推理失败: {result}")
                
        except Exception as e:
            self.logger.error(f"触发边侧节点推理异常: {node.node_id}, {str(e)}")
            raise
    
    async def get_system_status(self, node: Node) -> EdgeSystemStatus:
        """获取边侧节点系统状态"""
        try:
            client = self._get_http_client(node)
            success, result = client.get("status")
            
            if success:
                self.logger.debug(f"获取边侧节点系统状态成功: {node.node_id}")
                return EdgeSystemStatus(
                    node_id=node.node_id,
                    cpu_usage=result.get("cpu_usage", 0.0),
                    memory_usage=result.get("memory_usage", 0.0),
                    disk_usage=result.get("disk_usage", 0.0),
                    network_status=result.get("network_status", "unknown"),
                    service_status=result.get("service_status", {}),
                    timestamp=datetime.now().isoformat()
                )
            else:
                raise NetworkException(f"获取系统状态失败: {result}")
                
        except Exception as e:
            self.logger.error(f"获取边侧节点系统状态异常: {node.node_id}, {str(e)}")
            raise
    
    async def request_model_download(self, node: Node) -> EdgeOperationResult:
        """请求边侧节点从云端下载模型"""
        try:
            client = self._get_http_client(node)
            success, result = client.post("model/download")
            
            if success:
                self.logger.debug(f"请求边侧节点下载模型成功: {node.node_id}")
                return EdgeOperationResult(
                    success=True,
                    node_id=node.node_id,
                    operation="request_model_download",
                    data=result,
                    timestamp=datetime.now().isoformat()
                )
            else:
                raise NetworkException(f"请求模型下载失败: {result}")
                
        except Exception as e:
            self.logger.error(f"请求边侧节点下载模型异常: {node.node_id}, {str(e)}")
            raise
    
    async def get_available_models(self, node: Node) -> EdgeOperationResult:
        """获取边侧节点可用模型列表"""
        try:
            client = self._get_http_client(node)
            success, result = client.get("model/available")
            
            if success:
                self.logger.debug(f"获取边侧节点可用模型成功: {node.node_id}")
                return EdgeOperationResult(
                    success=True,
                    node_id=node.node_id,
                    operation="get_available_models",
                    data=result,
                    timestamp=datetime.now().isoformat()
                )
            else:
                raise NetworkException(f"获取可用模型失败: {result}")
                
        except Exception as e:
            self.logger.error(f"获取边侧节点可用模型异常: {node.node_id}, {str(e)}")
            raise
    
    async def check_cloud_model_status(self, node: Node) -> EdgeOperationResult:
        """检查边侧节点与云端模型状态"""
        try:
            client = self._get_http_client(node)
            success, result = client.get("model/cloud-status")
            
            if success:
                self.logger.debug(f"检查边侧节点云端模型状态成功: {node.node_id}")
                return EdgeOperationResult(
                    success=True,
                    node_id=node.node_id,
                    operation="check_cloud_model_status",
                    data=result,
                    timestamp=datetime.now().isoformat()
                )
            else:
                raise NetworkException(f"检查云端模型状态失败: {result}")
                
        except Exception as e:
            self.logger.error(f"检查边侧节点云端模型状态异常: {node.node_id}, {str(e)}")
            raise 